from config import InferenceConfig
from kvcache import KVCache
import numpy as np
from typing import List
import time
import sys


class Session:
    def __init__(self, config: InferenceConfig) -> None:
        self.kvCache = KVCache.create(config)
        self.run_times = 0
    def run(self, input_ids: np.ndarray):
        pass

    @staticmethod
    def fromConfig(config: InferenceConfig) -> 'Session':
        if config.session_type == "onnx":
            return OnnxSession(config)
        elif config.session_type == 'acl':
            return AclSession(config)
        else:
            return None

    def reset(self):
        if self.run_times == 0:
            self.kvCache.reset(0)
        else:
            self.kvCache.reset()

    def rollback(self, seq_len):
        self.kvCache.rollback(seq_len)


class OnnxSession(Session):
    def __init__(self, config: InferenceConfig) -> None:
        super().__init__(config)
        import onnxruntime
        options = onnxruntime.SessionOptions()
        self.llm_session = onnxruntime.InferenceSession(
            config.model,
            sess_options=options,
            providers=[
                "DmlExecutionProvider",
                "CUDAExecutionProvider",
                "CPUExecutionProvider",
            ],
        )

    def run(self, input_ids: np.ndarray):
        seq_len = input_ids.shape[-1]
        cache, mask, pos_ids = self.kvCache.getInputs(seq_len)
        result = self.llm_session.run(None, {
            "input_ids": input_ids,
            "attention_mask": mask,
            "past_key_values": cache,
            "position_ids": pos_ids,
        })
        self.kvCache.update(seq_len, result[1])
        return result


class AclSession(Session):
    context = None

    def __init__(self, config: InferenceConfig) -> None:
        super().__init__(config)
        from engine import ACLModel, initResource
        # self.context = initResource(config.device)
        self.model = ACLModel(config.model, self.context)
        self.input_embeds = np.zeros((1, 16, 2028), dtype=np.int64)
        self.run_times = 0
        self.kvCache.kvCache = self.model.kvCache
    def run(self, input_embeds: np.ndarray):
        seq_len, logits = input_embeds.shape[-2], None
        for i in range(seq_len):
            logits = self.run_one(input_embeds[:, i])
        return [logits]


    def run_all_logits(self, input_embeds: np.ndarray):
        seq_len, i = input_ids.shape[-1], 0
        logits = []
        while i < seq_len:
            end = i + 16 if i + 16 < seq_len else seq_len
            cache, mask, pos_ids = self.kvCache.getInputs(16)
            self.input_ids[0:end - i] = input_ids[i:end]
            result: List[np.ndarray] = self.model.inference([self.input_ids, mask, pos_ids, cache])
            self.kvCache.update(end - i, result[1])
            logits.append(result[0][0:end - i].reshape(1, -1))
        return [np.concatenate(logits).reshape(1, 1, -1)]


    def run_one(self, input_embeds: np.ndarray):
        self.run_times += 1
        start = time.time()     
        cache, mask, pos_ids = self.kvCache.getInputs(1)
        result: List[np.ndarray] = self.model.inference([mask, pos_ids, cache, input_embeds])
        self.kvCache.update(1, result[1])
        return result[0].reshape(1, 1, -1)
    def run_picture(self,input_embeds:np.ndarray):
        self.run_times += 576
        cache,mask,pos_ids = self.kvCache.getInputs(576)
        result:List[np.ndarray] = self.model.inference([mask, pos_ids, cache, input_embeds],True)
        self.kvCache.update(576,result[1])
        return result[0][-102400:].reshape(1,1,-1)
        
    def run_prefill(self, input_embeds: np.ndarray, images_seq_mask):
        seq_len,logits = input_embeds.shape[-2],None
        index = 0
        while index < seq_len:
            if images_seq_mask[0][index]:
                logits = self.run_picture(input_embeds[0][index:index+576])
                index = index + 576
            else:
                logits = self.run_one(input_embeds[0][index])
                index = index + 1
        return [logits]
